import arrow

from apps.event.serializer import EventActionSerializer
from apps.notify.display import Tips

from models.target import Targets
from models.task import ScanTasks, ScanTaskSessions, ScanEvents, ScanEventActions, ScanTaskSessionsExtra, ScanPages
from utilities.meta_view import (
    Response,
    APIViewList,
    APIViewUpdate,
    APIViewRetrieve,
    Success,
    APIViewCreate,
)
from utilities.scanner import BatchStopTask, VulVerificationScheduler
from utilities.utils import filter_cha
from utilities.log import audit_log
from utilities.enums import LogModuleName, AuditLogAction, LogTypeName, TaskType, TaskStatusName
from utilities.task_config import TASK_TYPE_TUPLE, EXTaskType
from utilities.djconsole import djconsole


class ListScanTask(APIViewList):
    check_user = True

    def list(self):
        page, limit = self.get_paginate()
        task_type = self.request.GET.get('task_type')
        status = self.request.GET.get('status')

        targets_id_map = self.get_target_id_map()

        task_query = {
            'target_id': {'$in': list(targets_id_map.keys())},
            'task_config.is_disabled': False,
            'addition.task_type': {'$ne': EXTaskType.target_info.value},
        }
        if status:
            status = int(status)
            if status == 1:
                task_query['status'] = {'$in': [1, 2]}
            else:
                task_query['status'] = status
        if task_type:
            task_query['addition.task_type'] = task_type

        all_tasks = (
            ScanTasks.objects.find(task_query)
            .order_by('-last_started_at')
            .paginate(limit)
        )
        tasks = all_tasks.page(page)

        data = []
        for task in tasks:
            task_id = task.task_id
            target = targets_id_map.get(task.target_id, {})
            temp_dict = {
                'task_id': task_id,
                'target_url': target.get('target_url', ''),
                'target_name': target.get('target_name', ''),
                'target_title': target.get('target_title', ''),
                'target_group_name': target.get('target_group', {}).get(
                    'target_group_name'
                ),
                'task_type': task.addition.get('task_type', ''),
                'status': task.status,
                'progress': task.progress,
                'is_disabled': task.task_config.get('is_disabled'),
                'task_session_id': task.last_task_session_id,
                'last_started_at': arrow.get(task.last_started_at).isoformat()
                if task.last_started_at
                else '',
                'next_start_at': arrow.get(task.next_start_at).isoformat()
                if task.next_start_at
                else '',
            }
            data.append(temp_dict)

        return Response({'data': data, 'page': page, 'total': all_tasks.count})

    def get_target_id_map(self):
        target_name = self.request.GET.get('target_name')
        target_group_name = self.request.GET.get('target_group_name')
        target_url = self.request.GET.get('target_url')

        target_query = {**self.user_query}
        ex_match = {}
        if target_url:
            target_query['target_url'] = {'$regex': filter_cha(target_url)}
        if target_name:
            target_name = filter_cha(target_name)
            target_query['$or'] = [
                {'target_name': {'$regex': target_name}},
                {'target_name': '', 'target_title': {'$regex': target_name}},
            ]
        if target_group_name:
            ex_match['target_group.target_group_name'] = {
                '$regex': filter_cha(target_group_name)
            }

        pipeline = [
            {'$match': target_query},
            {
                '$lookup': {
                    'from': "scan_target_groups",
                    'localField': "target_group_id",
                    'foreignField': "target_group_id",
                    'as': "target_group",
                }
            },
            {'$unwind': '$target_group'},
            {'$match': ex_match},
        ]

        targets_id_map = {
            target.get('target_id'): target
            for target in Targets.objects.aggregate(pipeline)
        }

        return targets_id_map


class CancelScanTask(APIViewUpdate):
    check_user = True

    def update(self):
        task_id_list = self.request.data.get('task_ids', [])
        query = {'task_id': {'$in': task_id_list}, **self.user_query}

        pipeline = [
            {'$match': query},
            {
                '$lookup': {
                    'from': "scan_targets",
                    'localField': "target_id",
                    'foreignField': "target_id",
                    'as': "target",
                }
            },
            {'$unwind': '$target'},
        ]
        task_id_list = []
        task_target_url_list = []
        task_session_id_list = []
        for task in ScanTasks.objects.aggregate(pipeline):
            task_id_list.append(task.get('task_id'))
            task_target_url_list.append(task.get('target_url'))
            task_session_id_list.append(task.get('last_task_session_id'))

        BatchStopTask(ids=task_id_list).batch_stop()
        ScanTaskSessionsExtra.objects.filter(task_session_id__in=task_session_id_list).update(
            **{'status': 8, 'reason': Tips().get(0), 'status_name': TaskStatusName.aborted.value})

        log_data = {
            "module": LogModuleName.task_management.value,
            "action": AuditLogAction.del_.value,
            "type": LogTypeName.task.value,
            "desc": {"detail": f"取消任务 {', '.join(task_target_url_list)}"},
        }
        audit_log(self.request, **log_data)

        return Success()


class ScanTaskSummary(APIViewRetrieve):
    check_user = True

    def retrieve(self):
        query = {
            'target_id': {'$ne': ''},
            'task_config.is_disabled': False,
            **self.user_query,
            'addition.task_type': {'$in': TASK_TYPE_TUPLE}
        }

        task_type_tuple_ne_ava = tuple(item for item in TASK_TYPE_TUPLE if item != TaskType.availability.value)
        queue_count = ScanTasks.objects.find(
            {'status': {'$in': [1, 2]}, **query}
        ).count()
        processing_count = ScanTasks.objects.find({'status': 3, **query}).count()
        task_session_count = ScanTaskSessions.objects.find(
            {'addition.task_type': {'$in': task_type_tuple_ne_ava}, **self.user_query}).count()
        target_count = tuple(
            ScanTasks.objects.aggregate(
                [
                    {'$match': query},
                    {'$group': {'_id': '$target_id'}},
                    {'$count': 'count'},
                ]
            )
        )

        url_count = ScanPages.objects.find(self.user_query).count()
        target_count = target_count[0].get('count', 0) if target_count else 0

        data = {
            'queue_count': queue_count,
            'processing_count': processing_count,
            'task_session_count': task_session_count,
            'target_count': target_count,
            'url_count': url_count,
        }

        return Response({'status': 'success', 'data': data})


class VulVerification(APIViewCreate):
    check_user = True

    def create(self):
        license_info = djconsole.get_license_detail(self.request)
        if EXTaskType.vul_verification_task.value not in license_info.get('modules', []):
            return Response({'status': 'failed', 'message': "Not Valid."})

        event_id = self.request.data.get('event_id')
        obj_event = ScanEvents.objects.find_one({'event_id': event_id, **self.user_query})
        obj_session = ScanTaskSessions.objects.find_one(
            {'task_session_id': obj_event.task_session_id, **self.user_query}
        )

        kwargs = {
            'event': obj_event.to_dict(),
            'task_session': obj_session.job_config_map,
        }

        VulVerificationScheduler(**kwargs).create()
        event_action = ScanEventActions.objects.filter(event_id=event_id).first()
        if event_action:
            event_action.update(verification=1)
        else:
            serializer = EventActionSerializer(
                data={'event_id': event_id, 'verification': 1, **obj_session.to_dict()}
            )
            serializer.is_valid(raise_exception=True)
            serializer.save()

        data = {
            "module": LogModuleName.task_management.value,
            "action": AuditLogAction.add.value,
            "type": LogTypeName.task.value,
            "desc": {"detail": f"验证web漏洞。"},
        }
        audit_log(request=self.request, **data)

        return Success()
